#include "masked_fill_custom_tiling.h"
#include "kernel_operator.h"

constexpr uint32_t BUFFER_NUM = 2;   // 双缓冲队列
constexpr uint32_t ALIGN_BYTES = 32; // 32字节对齐基准

constexpr uint32_t SIG_BFLOAT16 = 0;
constexpr uint32_t SIG_FLOAT16 = 1;
constexpr uint32_t SIG_FLOAT32 = 2;
constexpr uint32_t SIG_FLOAT64 = 3;
constexpr uint32_t SIG_UINT8 = 4;
constexpr uint32_t SIG_UINT16 = 5;
constexpr uint32_t SIG_UINT32 = 6;
constexpr uint32_t SIG_UINT64 = 7;
constexpr uint32_t SIG_INT8 = 8;
constexpr uint32_t SIG_INT16 = 9;
constexpr uint32_t SIG_INT32 = 10;
constexpr uint32_t SIG_INT64 = 11;
constexpr uint32_t DATA_TYPE_SIZE[] = {2, 2, 4, 8, 1, 2, 4, 8, 1, 2, 4, 8};

template <typename T>
class KernelMaskedFill
{
public:
    __aicore__ inline KernelMaskedFill() {}
    __aicore__ inline void Init(GM_ADDR x, GM_ADDR mask, GM_ADDR valueT, GM_ADDR z, MaskedFillCustomTilingData tiling)
    {
        const T maskValue = *reinterpret_cast<const __gm__ T*>(valueT);
        this->value = maskValue;
        this->xTypeSize = DATA_TYPE_SIZE[tiling.dataType];
        this->localValidElem = 0U;                    // 初始化局部有效元素数为0

        this->alignNum_x = ALIGN_BYTES / this->xTypeSize;         // x的对齐元素数
        this->alignNum_mask = ALIGN_BYTES / sizeof(bool);         // mask的对齐元素数

        uint32_t blockIdx = AscendC::GetBlockIdx();
        uint32_t gmStartIdx = 0U;
        uint32_t blockTotalElemAligned = 0U; 

        if (tiling.isEvenCore)
        {
            blockTotalElemAligned = tiling.blockLength;
            gmStartIdx = blockTotalElemAligned * blockIdx;
            this->tileNum = tiling.tileNum;
            this->tileLength_x = tiling.tileLength;
            this->tileLength_mask = AlignUp(tiling.tileLength, alignNum_mask);
            this->lastTileLength_x = tiling.lastTileLength;
            this->lastTileLength_mask = AlignUp(tiling.lastTileLength, alignNum_mask);
        }
        else
        {
            if (blockIdx < tiling.formerNum)
            {
                blockTotalElemAligned = tiling.formerLength;
                gmStartIdx = blockTotalElemAligned * blockIdx;
                this->tileNum = tiling.formerTileNum;
                this->tileLength_x = tiling.formerTileLength;
                this->tileLength_mask = AlignUp(tiling.formerTileLength, alignNum_mask);
                this->lastTileLength_x = tiling.formerLastTileLength;
                this->lastTileLength_mask = AlignUp(tiling.formerLastTileLength, alignNum_mask);
            }
            else
            {
                blockTotalElemAligned = tiling.tailLength;
                gmStartIdx = tiling.formerLength * tiling.formerNum + 
                             blockTotalElemAligned * (blockIdx - tiling.formerNum);
                this->tileNum = tiling.tailTileNum;
                this->tileLength_x = tiling.tailTileLength;
                this->tileLength_mask = AlignUp(tiling.tailTileLength, alignNum_mask);
                this->lastTileLength_x = tiling.tailLastTileLength;
                this->lastTileLength_mask = AlignUp(tiling.tailLastTileLength, alignNum_mask);
            }
        }

        this->gmStartIdx = gmStartIdx;
        if (this->gmStartIdx >= tiling.totalLength)
        {
            this->localValidElem = 0U;
            return;
        }
        this->localValidElem = (this->gmStartIdx + blockTotalElemAligned > tiling.totalLength) ? 
                               (tiling.totalLength - this->gmStartIdx) : blockTotalElemAligned;

        xGm.SetGlobalBuffer(reinterpret_cast<__gm__ T*>(x) + this->gmStartIdx, this->localValidElem);
        maskGm.SetGlobalBuffer(reinterpret_cast<__gm__ bool*>(mask) + this->gmStartIdx, this->localValidElem);
        zGm.SetGlobalBuffer(reinterpret_cast<__gm__ T*>(z) + this->gmStartIdx, this->localValidElem);

        uint32_t xBufSize = this->tileLength_x * this->xTypeSize;
        uint32_t maskBufSize = this->tileLength_mask * sizeof(bool);

        pipe.InitBuffer(inQueueX, BUFFER_NUM, xBufSize);
        pipe.InitBuffer(inQueueM, BUFFER_NUM, maskBufSize);
        pipe.InitBuffer(outQueueZ, BUFFER_NUM, xBufSize);
    }

    __aicore__ inline void Process()
    {
        if (this->localValidElem == 0U || this->tileNum == 0U) return;
        uint32_t loopCount = this->tileNum * BUFFER_NUM;
        for (uint32_t i = 0U; i < loopCount; i++)
        {
            CopyIn(i);
            Compute(i);
            CopyOut(i);
        }
    }

private:
    __aicore__ inline uint32_t AlignUp(uint32_t value, uint32_t align)
    {
        return (value == 0U) ? 0U : ((value + align - 1U) & ~(align - 1U));
    }

    __aicore__ inline uint32_t GetCopyCount(uint32_t progress, uint32_t lastTileLength, uint32_t tileLength, uint32_t alignNum)
    {
        uint32_t tileIdx = progress / BUFFER_NUM;
        uint32_t tileLen = (tileIdx == this->tileNum - 1U) ? lastTileLength : tileLength;
        uint32_t offsetInCore = tileIdx * this->tileLength_x; 
        uint32_t maxValid = this->localValidElem - offsetInCore;

        if (maxValid <= 0U) return 0U;
        return (tileLen <= maxValid) ? tileLen : AlignUp(maxValid, alignNum);
    }

    __aicore__ inline void CopyIn(uint32_t progress)
    {
        uint32_t tileIdx = progress / BUFFER_NUM;
        uint32_t offsetInCore = tileIdx * this->tileLength_x;
        if (offsetInCore >= this->localValidElem)
        {
            this->xCopyCountCache[progress % BUFFER_NUM] = 0U;
            this->maskCopyCountCache[progress % BUFFER_NUM] = 0U;
            return;
        }

        AscendC::LocalTensor<T> xLocal = inQueueX.AllocTensor<T>();
        AscendC::LocalTensor<bool> maskLocal = inQueueM.AllocTensor<bool>();

        uint32_t xCopyCount = GetCopyCount(progress, this->lastTileLength_x, this->tileLength_x, this->alignNum_x);
        uint32_t maskCopyCount = GetCopyCount(progress, this->lastTileLength_mask, this->tileLength_mask, this->alignNum_mask);
        if (xCopyCount > 0U) AscendC::DataCopy(xLocal, xGm[offsetInCore], xCopyCount);
        if (maskCopyCount > 0U) AscendC::DataCopy(maskLocal, maskGm[offsetInCore], maskCopyCount);

        this->xCopyCountCache[progress % BUFFER_NUM] = xCopyCount;
        this->maskCopyCountCache[progress % BUFFER_NUM] = maskCopyCount;
        inQueueX.EnQue(xLocal);
        inQueueM.EnQue(maskLocal);
    }

    __aicore__ inline void Compute(uint32_t progress)
    {
        uint32_t tileIdx = progress / BUFFER_NUM;
        uint32_t bufIdx = progress % BUFFER_NUM;
        uint32_t offsetInCore = tileIdx * this->tileLength_x;
        if (offsetInCore >= this->localValidElem) return;
        
        AscendC::LocalTensor<T> xLocal = inQueueX.DeQue<T>();
        AscendC::LocalTensor<bool> maskLocal = inQueueM.DeQue<bool>();
        AscendC::LocalTensor<T> zLocal = outQueueZ.AllocTensor<T>();
        uint32_t xCopyCount = this->xCopyCountCache[bufIdx];
        uint32_t maskCopyCount = this->maskCopyCountCache[bufIdx];
        uint32_t maxValid = this->localValidElem - offsetInCore;
        uint32_t result = AscendC::Std::min(xCopyCount, maskCopyCount); 
        uint32_t computeCount = AscendC::Std::min(result, maxValid);

        for (uint32_t i = 0U; i < computeCount; i++)
        {
            zLocal.SetValue(i, maskLocal.GetValue(i) ? this->value : xLocal.GetValue(i));
        }

        outQueueZ.EnQue(zLocal);
        inQueueX.FreeTensor(xLocal);
        inQueueM.FreeTensor(maskLocal);
    }

    __aicore__ inline void CopyOut(uint32_t progress)
    {
        uint32_t tileIdx = progress / BUFFER_NUM;
        uint32_t bufIdx = progress % BUFFER_NUM;
        uint32_t offsetInCore = tileIdx * this->tileLength_x;
        if (offsetInCore >= this->localValidElem) return;

        AscendC::LocalTensor<T> zLocal = outQueueZ.DeQue<T>();
        uint32_t copyCount = GetCopyCount(progress, this->lastTileLength_x, this->tileLength_x, this->alignNum_x);
        if (copyCount > 0U) AscendC::DataCopy(zGm[offsetInCore], zLocal, copyCount);

        outQueueZ.FreeTensor(zLocal);
    }

private:
    AscendC::TPipe pipe;
    AscendC::TQue<AscendC::TPosition::VECIN, BUFFER_NUM> inQueueX, inQueueM;
    AscendC::TQue<AscendC::TPosition::VECOUT, BUFFER_NUM> outQueueZ;

    AscendC::GlobalTensor<T> xGm;
    AscendC::GlobalTensor<bool> maskGm;
    AscendC::GlobalTensor<T> zGm;

    uint32_t alignNum_x;          // x的对齐元素数
    uint32_t alignNum_mask;       // mask的对齐元素数
    uint32_t tileNum;             // Tile数量
    uint32_t tileLength_x;        // x的Tile长度（对齐后）
    uint32_t tileLength_mask;     // mask的Tile长度（对齐后）
    uint32_t lastTileLength_x;    // x的尾Tile长度（对齐后）
    uint32_t lastTileLength_mask; // mask的尾Tile长度（对齐后）

    uint32_t gmStartIdx;          // GM起始索引
    uint32_t localValidElem;      // 局部有效元素数

    uint32_t xCopyCountCache[BUFFER_NUM] = {0U};
    uint32_t maskCopyCountCache[BUFFER_NUM] = {0U};

    uint32_t xTypeSize;
    T value;                      
};

// 核函数入口
extern "C" __global__ __aicore__ void masked_fill_custom(GM_ADDR x, 
                                                       GM_ADDR mask, 
                                                       GM_ADDR value, 
                                                       GM_ADDR z, 
                                                       MaskedFillCustomTilingData tiling)
{
    switch (tiling.dataType)
    {
        case SIG_BFLOAT16: { KernelMaskedFill<bfloat16_t> op; op.Init(x, mask, value, z, tiling); op.Process(); break; }
        case SIG_FLOAT16:  { KernelMaskedFill<half> op; op.Init(x, mask, value, z, tiling); op.Process(); break; }
        case SIG_FLOAT32:  { KernelMaskedFill<float> op; op.Init(x, mask, value, z, tiling); op.Process(); break; }
        case SIG_FLOAT64:  { KernelMaskedFill<double> op; op.Init(x, mask, value, z, tiling); op.Process(); break; }
        case SIG_UINT16:   { KernelMaskedFill<uint16_t> op; op.Init(x, mask, value, z, tiling); op.Process(); break; }
        case SIG_UINT32:   { KernelMaskedFill<uint32_t> op; op.Init(x, mask, value, z, tiling); op.Process(); break; }
        case SIG_UINT64:   { KernelMaskedFill<uint64_t> op; op.Init(x, mask, value, z, tiling); op.Process(); break; }
        case SIG_INT8:     { KernelMaskedFill<int8_t> op; op.Init(x, mask, value, z, tiling); op.Process(); break; }
        case SIG_INT16:    { KernelMaskedFill<int16_t> op; op.Init(x, mask, value, z, tiling); op.Process(); break; }
        case SIG_INT32:    { KernelMaskedFill<int> op; op.Init(x, mask, value, z, tiling); op.Process(); break; }
        case SIG_INT64:    { KernelMaskedFill<long> op; op.Init(x, mask, value, z, tiling); op.Process(); break; }
        default: assert(false && "unsupported dataType"); break;
    }
}